import numpy as np
import matplotlib.pyplot as plt
from .error_checker import MetricsErrorChecker

# 分类任务评估指标（多分类）

def accuracy(y_true, y_pred):
    """
    计算准确率
    """
    MetricsErrorChecker.check_is_array(y_true)
    MetricsErrorChecker.check_is_array(y_pred)
    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    MetricsErrorChecker.check_shapes_equal(y_true, y_pred)

    return np.sum(np.array(y_true) == np.array(y_pred)) / len(y_true)

def precision(y_true, y_pred, num_classes):
    """
    计算精确度（宏平均）
    """
    MetricsErrorChecker.check_valid_classification_inputs(y_true, y_pred, num_classes)

    precision_scores = []
    for class_id in range(num_classes):
        tp = np.sum((np.array(y_true) == class_id) & (np.array(y_pred) == class_id))  # 真阳性
        fp = np.sum((np.array(y_true) != class_id) & (np.array(y_pred) == class_id))  # 假阳性
        precision_scores.append(tp / (tp + fp) if tp + fp != 0 else 0)
    return np.mean(precision_scores)

def recall(y_true, y_pred, num_classes):
    """
    计算召回率（宏平均）
    """
    MetricsErrorChecker.check_valid_classification_inputs(y_true, y_pred, num_classes)

    recall_scores = []
    for class_id in range(num_classes):
        tp = np.sum((np.array(y_true) == class_id) & (np.array(y_pred) == class_id))  # 真阳性
        fn = np.sum((np.array(y_true) == class_id) & (np.array(y_pred) != class_id))  # 假阴性
        recall_scores.append(tp / (tp + fn) if tp + fn != 0 else 0)
    return np.mean(recall_scores)

def f1(y_true, y_pred, num_classes):
    """
    计算 F1 分数（宏平均）
    """
    MetricsErrorChecker.check_valid_classification_inputs(y_true, y_pred, num_classes)

    precision_scores = precision(y_true, y_pred, num_classes)
    recall_scores = recall(y_true, y_pred, num_classes)
    return 2 * (precision_scores * recall_scores) / (precision_scores + recall_scores) if (precision_scores + recall_scores) != 0 else 0

def confusion_matrix(y_true, y_pred, num_classes):
    """
    计算混淆矩阵（多分类）
    """
    MetricsErrorChecker.check_valid_classification_inputs(y_true, y_pred, num_classes)

    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    matrix = np.zeros((num_classes, num_classes), dtype=int)

    for true, pred in zip(y_true, y_pred):
        matrix[true, pred] += 1

    return matrix


# 回归任务评估指标

def mse(y_true, y_pred):
    """
    计算均方误差 (MSE)
    """
    MetricsErrorChecker.check_valid_regression_inputs(y_true, y_pred)

    return np.mean((np.array(y_true) - np.array(y_pred)) ** 2)

def mae(y_true, y_pred):
    """
    计算平均绝对误差 (MAE)
    """
    MetricsErrorChecker.check_valid_regression_inputs(y_true, y_pred)

    return np.mean(np.abs(np.array(y_true) - np.array(y_pred)))

def r2(y_true, y_pred):
    """
    计算 R² 分数
    """
    MetricsErrorChecker.check_valid_regression_inputs(y_true, y_pred)

    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    ss_total = np.sum((y_true - np.mean(y_true)) ** 2)
    ss_residual = np.sum((y_true - y_pred) ** 2)
    return 1 - (ss_residual / ss_total) if ss_total != 0 else 0  # 防止除以零

def rmse(y_true, y_pred):
    """
    计算均方根误差 (RMSE)
    """
    MetricsErrorChecker.check_valid_regression_inputs(y_true, y_pred)

    return np.sqrt(mse(y_true, y_pred))

def compute_roc(y_true, y_scores, num_classes):
    """
    手动计算多分类的 ROC 曲线和 AUC 值。
    """
    MetricsErrorChecker.check_valid_roc_inputs(y_true, y_scores, num_classes)

    y_true = np.array(y_true)
    fpr = {}
    tpr = {}
    aucs = {}

    for i in range(num_classes):
        # 针对每个类别，将其视为正类，其他类为负类
        y_true_bin = (y_true == i).astype(int)  # 将当前类别设为正类，其他类别为负类
        y_scores_bin = y_scores[:, i]  # 获取当前类别的预测概率

        # 计算 FPR 和 TPR
        thresholds = np.linspace(0, 1, 100)  # 取100个阈值
        tpr_values = []
        fpr_values = []
        for threshold in thresholds:
            y_pred_bin = (y_scores_bin >= threshold).astype(int)  # 按照阈值分类
            tp = np.sum((y_true_bin == 1) & (y_pred_bin == 1))  # 真阳性
            fp = np.sum((y_true_bin == 0) & (y_pred_bin == 1))  # 假阳性
            fn = np.sum((y_true_bin == 1) & (y_pred_bin == 0))  # 假阴性
            tn = np.sum((y_true_bin == 0) & (y_pred_bin == 0))  # 真阴性

            tpr_values.append(tp / (tp + fn) if tp + fn != 0 else 0)  # TPR
            fpr_values.append(fp / (fp + tn) if fp + tn != 0 else 0)  # FPR

        fpr[i] = np.array(fpr_values)
        tpr[i] = np.array(tpr_values)

        # 计算 AUC（使用梯形法则）
        aucs[i] = np.trapz(tpr[i], fpr[i])  # 用梯形法则计算 AUC

    return fpr, tpr, aucs


def plot_roc(fpr, tpr, aucs, num_classes):
    """
    绘制多分类的 ROC 曲线。
    """
    plt.figure()

    for i in range(num_classes):
        plt.plot(fpr[i], tpr[i], lw=2, label=f'Class {i} (AUC = {aucs[i]:.2f})')

    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')  # 随机猜测线
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic (ROC) Curve (Multi-class)')
    plt.legend(loc='lower right')
    plt.show()

def roc(y_true, y_scores, num_classes):
    """
    计算并绘制多分类的 ROC 曲线。
    """
    fpr, tpr, aucs = compute_roc(y_true, y_scores, num_classes)
    plot_roc(fpr, tpr, aucs, num_classes)
    return fpr, tpr, aucs
